import { tool } from '@langchain/core/tools';
import { z } from 'zod';
import { WeaviateService } from 'src/core/weaviate/weaviate.service';

const schema = z.object({
  query: z.string().describe('The search query'),
  n: z.number().optional().default(10).describe('Number of results to return'),
});

type Schema = z.infer<typeof schema>;

export const createSearchTool = (weaviateService: WeaviateService) => {
  const ragTool = tool(
    async ({ query, n = 10 }: Schema) => {
      console.log('查询工具调用', n, query);

      const results = await weaviateService.queryDocuments(query, n);

      const formattedResults = results
        .map(
          ([doc, score], index) =>
            `[${index + 1}] 相关度: ${score.toFixed(2)}\n${
              doc.pageContent
            }\n来源: ${doc.metadata.source || '未知'}\n`,
        )
        .join('\n');

      return `在知识库中找到以下相关内容:\n${formattedResults}`;
    },
    {
      name: 'hmi 2d图纸和hmi 3d场景信息检索工具',
      description: '信息检索工具',
      schema,
    },
  );

  return ragTool;
};
